package com.mentalhealthplatform.studentV3.ai.message;

import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.Resource;
import lombok.AllArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.model.Media;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.zhipuai.ZhiPuAiChatModel;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.util.CollectionUtils;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.List;

@RequestMapping("api/v3/message")
@RestController
@AllArgsConstructor
@Slf4j
public class AiMessageController {


    private final AiMessageRepository aiMessageRepository;
    private final AiMessageChatMemory chatMemory;
    private final ZhiPuAiChatModel zhiPuAiChatModel;
    private final ObjectMapper objectMapper;

    @Resource
    private OllamaChatModel ollamaChatModel;

    @DeleteMapping("history/{sessionId}")
    public void deleteHistory(@PathVariable String sessionId) {
        aiMessageRepository.deleteBySessionId(sessionId);
    }

    @GetMapping("history/{sessionId}")
    public List<AiMessageImpl> getHistory(@PathVariable String sessionId) {
        return aiMessageRepository.findBySessionId(sessionId);
    }

    @PostMapping
    public void save(@RequestBody AiMessageInput input) {
        AiMessageImpl aiMessage = input.toEntity();
        aiMessageRepository.save(aiMessage);
    }

    @PostMapping(value = "chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> chatStreamWithHistory(@RequestBody AiMessageInput input) {
        var messageChatMemoryAdvisor = new MessageChatMemoryAdvisor(chatMemory, input.getSessionId(), 10);
        return ChatClient.create(ollamaChatModel).prompt()
                .user(promptUserSpec -> {
                    Message message = (Message) AiMessageChatMemory.toSpringAiMessage(input.toEntity());
                    if (message instanceof UserMessage userMessage &&
                            !CollectionUtils.isEmpty(userMessage.getMedia())) {
                        Media[] medias = new Media[userMessage.getMedia().size()];
                        promptUserSpec.media(userMessage.getMedia().toArray(medias));
                    }
                    promptUserSpec.text(message.getContent());
                })
                .advisors(messageChatMemoryAdvisor)
                .stream()
                .content()
                .map(chatResponse -> ServerSentEvent.builder(toJson(chatResponse))
                        .event("message")
                        .build());
    }

    @PostMapping(value = "test/chat", produces = MediaType.APPLICATION_JSON_VALUE)
    public String chatWithHistory(@RequestBody AiMessageInput input) {
        var messageChatMemoryAdvisor = new MessageChatMemoryAdvisor(chatMemory, input.getSessionId(), 10);
        return ChatClient.create(zhiPuAiChatModel).prompt()
                .user(promptUserSpec -> {
                    Message message = (Message) AiMessageChatMemory.toSpringAiMessage(input.toEntity());
                    if (message instanceof UserMessage userMessage &&
                            !CollectionUtils.isEmpty(userMessage.getMedia())) {
                        Media[] medias = new Media[userMessage.getMedia().size()];
                        promptUserSpec.media(userMessage.getMedia().toArray(medias));
                    }
                    promptUserSpec.text(message.getContent());
                })
                .advisors(messageChatMemoryAdvisor)
                .call()
                .content();
    }


    @SneakyThrows
    public String toJson(String response) {
        return objectMapper.writeValueAsString(response);
    }
}